import javassist.bytecode.*;

import java.io.*;
import java.util.*;

import static javassist.bytecode.Opcode.*;

public class JavaSsist {
    private static ClassFile classFile;
    //inputFilePath={inputClassFileName}
    private static String inputClassFileName = "src/main/java/T.class";
    //outputFilePath={outputClassFileName}
    private static String outputClassFileName = "defaultOut.class";
    private static LocalVariableAttribute localVariableAttribute;
    private static int localVariableTableLength;

    public static void main(String[] args) {
        //read .class file
        init(args);
        try {
            BufferedInputStream fin = new BufferedInputStream(new FileInputStream(inputClassFileName));
            classFile = new ClassFile(new DataInputStream(fin));
            System.out.println("read file finish: " + inputClassFileName);
        } catch (IOException e) {
            e.printStackTrace();
            return;
        }

        //iterate methods in classfile
        List<MethodInfo> list = classFile.getMethods();
        for (MethodInfo method : list) {
            // 获取方法的本地变量表，如果没有使用javac -g 编译java文件，不会产生变量表
            CodeAttribute codeAttribute = (CodeAttribute) method.getAttribute(CodeAttribute.tag);
            localVariableAttribute = (LocalVariableAttribute) codeAttribute.getAttribute(LocalVariableAttribute.tag);
            localVariableTableLength = localVariableAttribute.tableLength();
            try {
                instrument(method); // instrument
            } catch (BadBytecode badBytecode) {
                badBytecode.printStackTrace();
            }

        }
        createFile(outputClassFileName);
        // output classfile
        try {
            classFile.write(new DataOutputStream(new FileOutputStream(outputClassFileName)));
            System.out.println("output file finish: " + outputClassFileName);

        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    private static void createFile(String filePath) {
        File file = new File(filePath);
        if (!file.exists()) {
            File pf = file.getParentFile();
            if (pf != null) {
                pf.mkdirs();
            }
            try {
                file.createNewFile();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private static void init(String[] args) {
        for (String str : args) {

            if (str.startsWith("inputFilePath=")) {
                inputClassFileName = str.substring(14);
                continue;
            }

            if (str.startsWith("outputFilePath=")) {
                outputClassFileName = str.substring(15);
            }
        }
    }

    private static void instrument(MethodInfo method) throws BadBytecode {
        CodeAttribute ca = method.getCodeAttribute();

        ca.setMaxStack(ca.getMaxStack() + 3); // To prevent insufficient stack depth.
        // print method  bytecode
        CodeIterator ci = ca.iterator();
        while (ci.hasNext()) {
            int index = ci.next(); // the position of next opcode
            int op = ci.byteAt(index); // get the value of opcode
//            System.out.println("\t" + index + " " + Mnemonic.OPCODE[op]);

            switch (op) {
                case INVOKESPECIAL: {
                    int num = ci.u16bitAt(index + 1);
                    try {
                        String methodName = findMethodNameInConstPool(classFile, num);
                        if (methodName.equals("<init>")) {
                            if (index - 4 >= 0 && (ci.byteAt(index - 4) == NEW)) {
                                insertNativeC(ci, index, 3, method, "newObj");
                            }
                        }
                    } catch (MethodNotFound methodNotFound) {
                        methodNotFound.printStackTrace();
                    }
                    break;
                }
                /**
                 * multianewarray
                 * indexbyte1
                 * indexbyte2
                 * dimensions
                 *
                 * anewarray
                 * indexbyte1
                 * indexbyte2
                 *
                 * newarray
                 * atype
                 */
                case NEWARRAY:
                case MULTIANEWARRAY:
                case ANEWARRAY: {
                    int step = 0;
                    switch (op) {
                        case NEWARRAY:
                            step = 2;
                            break;
                        case MULTIANEWARRAY:
                            step = 4;
                            break;
                        case ANEWARRAY:
                            step = 3;
                            break;
                    }
                    insertNativeC(ci, index, step, method, "newArr");
                    break;
                }
            }

        }
    }

    /**
     * @param ci         字节码迭代器
     * @param index      INVOKESPECIAL NEWARRAY MULTIANEWARRAY ANEWARRAY 指令的索引
     * @param step       index+step的位置为插入代码的位置
     * @param method     需要增强的方法
     * @param methodName 增强的方法的名称
     * @throws BadBytecode
     */
    private static void insertNativeC(CodeIterator ci, int index, int step, MethodInfo method, String methodName) throws BadBytecode {
        Bytecode bytecode = new Bytecode(classFile.getConstPool());
        int positionInLocalVariableTable = stackInObj(ci, index, step, bytecode);
//                            System.out.println(method.getName() + ": " + method.getLineNumber(index));
        String objName;
        if (localVariableAttribute != null && positionInLocalVariableTable > 0 && positionInLocalVariableTable < localVariableAttribute.tableLength()) {
            int len = localVariableAttribute.tableLength();
            int i = 0;
            for (; i < len; i++) {
                if(localVariableAttribute.index(i) == positionInLocalVariableTable){
                    break;
                }
            }
            objName = localVariableAttribute.variableName(i);
        } else {
            objName = "withoutName";
        }
        bytecode.addOpcode(DUP);
        bytecode.addInvokevirtual("java.lang.Object", "hashCode", "()I");
        bytecode.addOpcode(BIPUSH);
        bytecode.add(method.getLineNumber(index));
        bytecode.addLdc(objName);
        bytecode.addInvokestatic("java.lang.Thread", "currentThread", "()Ljava/lang/Thread;");
        bytecode.addInvokestatic("NativeC", methodName, "(Ljava/lang/Object;IILjava/lang/String;Ljava/lang/Thread;)V");
        byte[] bytes = bytecode.get();
//                    System.out.println(Arrays.toString(bytes));
        ci.insert(bytes);
//                    System.out.println(method.getName() + ": " + method.getLineNumber(index));
    }

    /**
     * 将初始化的实例对象压入栈
     *
     * @param ci
     * @param index
     * @param step
     * @param bytecode
     * @throws BadBytecode
     */
    private static int stackInObj(CodeIterator ci, int index, int step, Bytecode bytecode) throws BadBytecode {
        int position = isAnonymousInit(ci, index, step);
        if (position >= 0) {
//            System.out.println("index: " + position);
            ci.next();//skip astore
            bytecode.addAload(position);
            return position;
//            System.out.println(localVariableAttribute.variableName(position));
        } else {
            bytecode.addOpcode(Opcode.DUP);
            return -1;
        }
    }

    private static boolean isAstore_n(int bytecode) {
        return (bytecode == ASTORE_0 || bytecode == ASTORE_1 || bytecode == ASTORE_2 || bytecode == ASTORE_3);
    }

    /**
     * @param ci         code iterator of method
     * @param startIndex the position of invokespecial opcode
     * @return the position in LocalVariableTable, if it's anonymous return -1
     */
    private static int isAnonymousInit(CodeIterator ci, int startIndex, int step) {
        int bytecode = -1;
        if (ci.hasNext()) {
            bytecode = ci.byteAt(startIndex + step);
//            System.out.println("step: "+step +" bytecode: " + bytecode);
            if (isAstore_n(bytecode)) { // use astore_n opcode, the instance is not anonymous
                return bytecode - 75;
            }
            if (bytecode == ASTORE) { // return the index after Astore opcode
                return ci.byteAt(startIndex + step + 1);
            } else {
                return -1;
            }
        }
        //  ci has not next
        return -1;
    }

    private static String findMethodNameInConstPool(ClassFile classFile, int methodNumber) throws MethodNotFound {
        ConstPool constPool = classFile.getConstPool();
        String name = constPool.getMethodrefName(methodNumber);
        if (name == null) {
            throw new MethodNotFound("");
        }
        return name;
    }

    private static class MethodNotFound extends Exception {
        public MethodNotFound(String msg) {
            super(msg);
        }
    }
}
